# Modified from https://github.com/THUDM/HGB/blob/master/NC/benchmark/scripts/data_loader.py, aiming to load hgb's datasets.

import os
import numpy as np
import scipy.sparse as sp
from collections import Counter, defaultdict
from sklearn.metrics import f1_score
import time


class data_loader:
    def __init__(self, path):
        # 初始化数据加载器类，检查路径是否存在，如果是压缩文件则解压。如果路径不存在则抛出错误提示。然后加载节点、边和标签数据。
        self.path = path
        if os.path.exists(path):
            pass
        elif os.path.exists(path+'.zip'):
            os.system("unzip {} -d {}".format(path+'.zip', os.path.join(*path.split('/')[:-1])))
        else:
            assert False, ('Warning: HGB node classification datasets not downloaded'
                ', please download them from HGB repository `https://github.com/THUDM/HGB`'
                ' and put them under `./data/`')

        self.nodes = self.load_nodes()
        self.links = self.load_links()
        self.labels_train = self.load_labels('label.dat')
        self.labels_test = self.load_labels('label.dat.test')

    # 函数 get_sub_graph 用于从原始异构图中提取一个子图，只保留指定类型的节点。这个函数会进行就地更新，并返回旧的节点类型和边类型到新的节点类型和边类型的映射。
    # node_types_tokeep 是一个列表或集合，包含需要保留的节点类型。
    # 这个函数只支持整个类型的子图。（包含对应节点类型的子图）
    def get_sub_graph(self, node_types_tokeep):
        """
        node_types_tokeep is a list or set of node types that you want to keep in the sub-graph
        We only support whole type sub-graph for now.
        This is an in-place update function!
        return: old node type id to new node type id dict, old edge type id to new edge type id dict
        """
        # 将 node_types_tokeep 转换为集合 keep，方便后续判断某个节点类型是否需要保留。
        keep = set(node_types_tokeep)
        # 初始化新图的数据结构，包括新的节点、边、训练标签、测试标签等。
        # new_node_type 和 new_node_id 用于跟踪新图中节点类型和节点的编号。
        new_node_type = 0
        new_node_id = 0
        new_nodes = {'total':0, 'count':Counter(), 'attr':{}, 'shift':{}}
        new_links = {'total':0, 'count':Counter(), 'meta':{}, 'data':defaultdict(list)}
        new_labels_train = {'num_classes':0, 'total':0, 'count':Counter(), 'data':None, 'mask':None}
        new_labels_test = {'num_classes':0, 'total':0, 'count':Counter(), 'data':None, 'mask':None}
        # old_nt2new_nt 用于存储旧节点类型到新节点类型的映射。
        old_nt2new_nt = {}
        # old_idx 用于存储旧节点索引，以便后续提取数据。
        old_idx = []
        # 遍历所有节点类型，如果节点类型在 keep 中：
        for node_type in self.nodes['count']:
            if node_type in keep:
                # 将旧节点类型映射到新节点类型。
                nt = node_type
                nnt = new_node_type
                old_nt2new_nt[nt] = nnt
                cnt = self.nodes['count'][nt]
                # 更新新图中的节点总数、节点类型计数、节点属性和节点编号偏移。
                new_nodes['total'] += cnt
                new_nodes['count'][nnt] = cnt
                new_nodes['attr'][nnt] = self.nodes['attr'][nt]
                new_nodes['shift'][nnt] = new_node_id
                beg = self.nodes['shift'][nt]
                old_idx.extend(range(beg, beg+cnt))

                # 更新新图中的训练和测试标签计数。
                cnt_label_train = self.labels_train['count'][nt]
                new_labels_train['count'][nnt] = cnt_label_train
                new_labels_train['total'] += cnt_label_train
                cnt_label_test = self.labels_test['count'][nt]
                new_labels_test['count'][nnt] = cnt_label_test
                new_labels_test['total'] += cnt_label_test

                # 更新新的节点类型编号和节点编号。
                new_node_type += 1
                new_node_id += cnt

        # 将原始图中的训练和测试标签的类别数复制到新图。
        new_labels_train['num_classes'] = self.labels_train['num_classes']
        new_labels_test['num_classes'] = self.labels_test['num_classes']
        # 根据 old_idx 提取相应的训练和测试标签数据。
        for k in ['data', 'mask']:
            new_labels_train[k] = self.labels_train[k][old_idx]
            new_labels_test[k] = self.labels_test[k][old_idx]

        # 初始化旧边类型到新边类型的映射。
        old_et2new_et = {}
        new_edge_type = 0
        # 遍历所有边类型，如果边类型的起始和终止节点类型都在 keep 中：
        for edge_type in self.links['count']:
            h, t = self.links['meta'][edge_type]
            if h in keep and t in keep:
                # 将旧边类型映射到新边类型。
                et = edge_type
                net = new_edge_type
                old_et2new_et[et] = net
                # 更新新图中的边总数、边类型计数和边类型元数据。
                new_links['total'] += self.links['count'][et]
                new_links['count'][net] = self.links['count'][et]
                new_links['meta'][net] = tuple(map(lambda x:old_nt2new_nt[x], self.links['meta'][et]))
                # 根据 old_idx 提取相应的边数据。
                new_links['data'][net] = self.links['data'][et][old_idx][:, old_idx]
                new_edge_type += 1

        # 将原始图的数据结构更新为新图的数据结构。
        self.nodes = new_nodes
        self.links = new_links
        self.labels_train = new_labels_train
        self.labels_test = new_labels_test
        # 返回旧节点类型到新节点类型的映射和旧边类型到新边类型的映射。
        return old_nt2new_nt, old_et2new_et

    # 函数 get_meta_path 用于获取异构图中指定元路径（meta-path）的稀疏矩阵表示。
    # meta 是一个包含边类型的列表，表示要构造的元路径。
    # 函数返回一个稀疏矩阵，形状为 [node_num, node_num]，表示节点之间通过元路径相连的情况。
    def get_meta_path(self, meta=[]):
        """
        Get meta path matrix
            meta is a list of edge types (also can be denoted by a pair of node types)
            return a sparse matrix with shape [node_num, node_num]
        """
        # sp.eye 创建一个大小为 [node_num, node_num] 的单位矩阵，其中 node_num 是图中节点的总数。
        # 单位矩阵的作用是作为矩阵乘法的初始矩阵，表示从节点自己出发的路径。
        ini = sp.eye(self.nodes['total'])
        # 将 meta 中的每个元素通过 self.get_edge_type(x) 函数转换成边类型的索引。如果 meta 中的元素是节点类型对，则转换成相应的边类型索引。
        meta = [self.get_edge_type(x) for x in meta]
        # 遍历 meta 中的每个边类型 x。
        # self.links['data'][x] 表示边类型 x 的邻接矩阵。
        # 如果 x 是正数，则直接使用 self.links['data'][x] 进行矩阵乘法。
        # 如果 x 是负数，则使用 self.links['data'][-x - 1].T，表示反向边的邻接矩阵的转置进行矩阵乘法。
        # 将结果累积到 ini 中，逐步构造出元路径的稀疏矩阵表示。
        for x in meta:
            ini = ini.dot(self.links['data'][x]) if x >= 0 else ini.dot(self.links['data'][-x - 1].T)
        # 返回累积得到的稀疏矩阵 ini，该矩阵表示通过元路径从一个节点可以到达另一个节点的情况。
        return ini

    # dfs 函数是深度优先搜索，用于在元路径中查找路径。
    # 用于在异构图中根据元路径（meta-path）找到所有可能的路径。具体来说，该函数沿着指定的元路径，从当前节点出发，递归地查找所有可能的路径，并将它们存储在 meta_dict 中。
    # now：当前路径，表示从起始节点到当前节点的路径。
    # meta：剩余的元路径（即需要继续探索的边类型序列）。
    # meta_dict：存储找到的所有路径，键是起始节点，值是从该节点出发的所有符合元路径的路径。
    def dfs(self, now, meta, meta_dict):
        # 如果 meta 为空，表示元路径已经完全匹配，将当前路径 now 添加到 meta_dict 中，然后返回。
        # now[0] 是路径的起始节点，now 是完整路径。
        if len(meta) == 0:
            meta_dict[now[0]].append(now)
            return
        # 获得对应边类型的稀疏矩阵
        # 类似于这样links['data'][0]:
        # self.links['data'] = {
        #     0: sp.csr_matrix([[0, 1, 0], [0, 0, 1], [0, 0, 0]]),  # 边类型 0: 节点 0 -> 1 -> 2
        #     1: sp.csr_matrix([[0, 0, 0], [1, 0, 0], [0, 1, 0]])   # 边类型 1: 节点 1 -> 0 -> 1
        # }
        # meta[0] 是当前需要匹配的边类型。
        # 如果 meta[0] 是正数，使用 self.links['data'][meta[0]] 获取正向边的邻接矩阵。
        # 如果 meta[0] 是负数，使用 self.links['data'][-meta[0] - 1].T 获取反向边的邻接矩阵的转置。
        th_mat = self.links['data'][meta[0]] if meta[0] >= 0 else self.links['data'][-meta[0] - 1].T
        # now[-1] 是当前路径的最后一个节点，即当前节点。
        th_node = now[-1]
        # th_mat[th_node].nonzero()[1] 获取当前节点在邻接矩阵中的所有非零列索引，即所有邻居节点。
        # 对每个邻居节点 col，递归调用 dfs，将当前节点 col 添加到路径 now，并将剩余的元路径 meta[1:] 继续传递下去。
        for col in th_mat[th_node].nonzero()[1]:
            self.dfs(now+[col], meta[1:], meta_dict)

    # get_full_meta_path 函数返回完整的元路径，支持对称和非对称模式。
    # 该函数用于获取每个节点的完整元路径。
    # meta 是一个边类型的列表（也可以用节点类型对表示）。
    # 返回值是一个字典，其中键是节点 ID，值是节点 ID 列表的列表，表示元路径。
    def get_full_meta_path(self, meta=[], symmetric=False):
        """
        Get full meta path for each node
            meta is a list of edge types (also can be denoted by a pair of node types)
            return a dict of list[list] (key is node_id)
        """
        # 将 meta 中的每个元素通过 get_edge_type 转换为边类型 ID。
        meta = [self.get_edge_type(x) for x in meta]
        # 如果 meta 长度为 1，表示只需考虑一种边类型：
        if len(meta) == 1:
            # 初始化一个空字典 meta_dict。
            meta_dict = {}
            # 获取起始节点类型（从边类型元数据中获取，如果边类型 ID 为负数，则取反后获取）。
            # meta[0] 表示元路径中的第一个边类型。
            # self.links['meta'][meta[0]] 获取该边类型对应的节点类型对（起点类型和终点类型）。
            # 如果 meta[0] 是正数，说明这是一个正向边，则 self.links['meta'][meta[0]][0] 获取的是边的起点类型。
            # 如果 meta[0] 是负数，说明这是一个反向边，则 self.links['meta'][-meta[0] - 1][1] 获取的是边的终点类型。
            start_node_type = self.links['meta'][meta[0]][0] if meta[0]>=0 else self.links['meta'][-meta[0]-1][1]
            # 遍历所有属于起始节点类型的节点 ID。
            for i in range(self.nodes['shift'][start_node_type], self.nodes['shift'][start_node_type]+self.nodes['count'][start_node_type]):
                meta_dict[i] = []
                # 对每个节点，调用 dfs 函数进行深度优先搜索，生成元路径。
                self.dfs([i], meta, meta_dict)
        # 如果 meta 长度大于 1：
        else:
            # 初始化两个空字典 meta_dict1 和 meta_dict2。
            meta_dict1 = {}
            meta_dict2 = {}
            mid = len(meta) // 2
            # 将 meta 分成两部分：meta1 和 meta2。
            meta1 = meta[:mid]
            meta2 = meta[mid:]
            # 获取 meta1 的起始节点类型。
            start_node_type = self.links['meta'][meta1[0]][0] if meta1[0]>=0 else self.links['meta'][-meta1[0]-1][1]
            # 遍历所有属于起始节点类型的节点 ID。
            for i in range(self.nodes['shift'][start_node_type], self.nodes['shift'][start_node_type]+self.nodes['count'][start_node_type]):
                meta_dict1[i] = []
                # 对每个节点，调用 dfs 函数进行深度优先搜索，生成元路径并存储在 meta_dict1 中。
                self.dfs([i], meta1, meta_dict1)
            # 获取 meta2 的起始节点类型。
            start_node_type = self.links['meta'][meta2[0]][0] if meta2[0]>=0 else self.links['meta'][-meta2[0]-1][1]
            # 遍历所有属于起始节点类型的节点 ID，初始化 meta_dict2 中对应的键。
            for i in range(self.nodes['shift'][start_node_type], self.nodes['shift'][start_node_type]+self.nodes['count'][start_node_type]):
                meta_dict2[i] = []
            # 如果 symmetric 为真，则生成对称路径：
            # 遍历 meta_dict1 中的每个键（起始节点 ID）和对应的路径列表 paths。
            # 对于每条路径 x，将其反转，并将反转后的路径存储在 meta_dict2 中，以路径的终点节点为键。
            # 如果 symmetric 为假，则对每个起始节点类型的节点，调用 dfs 函数进行深度优先搜索，生成元路径并存储在 meta_dict2 中。
            if symmetric:
                for k in meta_dict1:
                    paths = meta_dict1[k]
                    for x in paths:
                        meta_dict2[x[-1]].append(list(reversed(x)))
            else:
                for i in range(self.nodes['shift'][start_node_type], self.nodes['shift'][start_node_type]+self.nodes['count'][start_node_type]):
                    self.dfs([i], meta2, meta_dict2)
            # 初始化一个空字典 meta_dict。
            meta_dict = {}
            # 获取 meta1 的起始节点类型。
            start_node_type = self.links['meta'][meta1[0]][0] if meta1[0]>=0 else self.links['meta'][-meta1[0]-1][1]
            # 对于每个起始节点，将 meta_dict1 和 meta_dict2 中的路径组合起来，形成完整的元路径：
            for i in range(self.nodes['shift'][start_node_type], self.nodes['shift'][start_node_type]+self.nodes['count'][start_node_type]):
                meta_dict[i] = []
                # 遍历 meta_dict1 中的起始路径 beg。
                for beg in meta_dict1[i]:
                    # 对于 beg 的终点节点，找到 meta_dict2 中对应的路径 end。
                    for end in meta_dict2[beg[-1]]:
                        # 将 beg 和 end（去掉起点部分）拼接起来，形成完整路径，存储在 meta_dict 中。
                        meta_dict[i].append(beg + end[1:])
        # 返回生成的 meta_dict，其中每个键是节点 ID，值是元路径列表。
        return meta_dict

    # gen_file_for_evaluate 函数生成用于评估的文件，支持二分类和多分类模式。
    def gen_file_for_evaluate(self, test_idx, label, file_name, mode='bi'):
        if test_idx.shape[0] != label.shape[0]:
            return
        if mode == 'multi':
            multi_label=[]
            for i in range(label.shape[0]):
                label_list = [str(j) for j in range(label[i].shape[0]) if label[i][j]==1]
                multi_label.append(','.join(label_list))
            label=multi_label
        elif mode=='bi':
            label = np.array(label)
        else:
            return
        with open(file_name, "w") as f:
            for nid, l in zip(test_idx, label):
                f.write(f"{nid}\t\t{self.get_node_type(nid)}\t{l}\n")

    # evaluate 函数根据预测结果计算微平均和宏平均的 F1 分数。
    def evaluate(self, pred):
        y_true = self.labels_test['data'][self.labels_test['mask']]
        micro = f1_score(y_true, pred, average='micro')
        macro = f1_score(y_true, pred, average='macro')
        result = {
            'micro-f1': micro,
            'macro-f1': macro
        }
        return result

    # 数据格式类似于下面这种：
    # 0	node0	0	0,1
    # 1	node1	1	2
    # 2	node2	0	1
    # 3	node3	1	0,2
    def load_labels(self, name):
        """
        return labels dict
            num_classes: total number of labels
            total: total number of labeled data
            count: number of labeled data for each node type
            data: a numpy matrix with shape (self.nodes['total'], self.labels['num_classes'])
            mask: to indicate if that node is labeled, if False, that line of data is masked
        """
        # 初始化一个 labels 字典，用于存储标签信息：
        # num_classes: 标签的总数，初始值为 0。
        # total: 有标签的数据总数，初始值为 0。
        # count: 每种节点类型的标签数量，用 Counter 统计。
        # data: 存储标签数据的矩阵，初始值为 None。
        # mask: 掩码矩阵，指示节点是否有标签，初始值为 None。
        labels = {'num_classes':0, 'total':0, 'count':Counter(), 'data':None, 'mask':None}
        # nc: 临时变量，用于计算标签的总数。
        # mask: 初始化一个大小为 self.nodes['total'] 的布尔数组，全部设为 False，表示初始时所有节点都没有标签。
        # data: 初始化一个大小为 self.nodes['total'] 的列表，全部设为 None，用于存储每个节点的标签。
        nc = 0
        mask = np.zeros(self.nodes['total'], dtype=bool)
        data = [None for i in range(self.nodes['total'])]
        # 打开标签文件（例如 label.dat），逐行读取数据。
        with open(os.path.join(self.path, name), 'r', encoding='utf-8') as f:
            for line in f:
                # 将每行数据按制表符 \t 分割成一个列表 th。
                th = line.split('\t')
                # 从列表中提取节点 ID (node_id)、节点名称 (node_name)、节点类型 (node_type) 和节点标签 (node_label)。
                # 将标签字符串转换为整数列表。
                node_id, node_name, node_type, node_label = int(th[0]), th[1], int(th[2]), list(map(int, th[3].split(',')))
                # 遍历 node_label 列表，更新 nc 为标签的最大值加 1，即计算标签的总数。
                for label in node_label:
                    nc = max(nc, label+1)
                # 将 mask 数组中对应 node_id 的位置设为 True，表示该节点有标签。
                mask[node_id] = True
                # 将 node_label 列表存储在 data 列表中对应 node_id 的位置。
                data[node_id] = node_label
                # 更新 labels['count'][node_type]，增加 1。
                labels['count'][node_type] += 1
                # 更新 labels['total']，增加 1。
                labels['total'] += 1
        # 设置 labels['num_classes'] 为 nc。
        labels['num_classes'] = nc
        # 初始化一个大小为 (self.nodes['total'], labels['num_classes']) 的零矩阵 new_data，用于存储标签数据。
        new_data = np.zeros((self.nodes['total'], labels['num_classes']), dtype=int)
        # 遍历 data 列表，如果对应位置有标签数据，将其转换为矩阵形式：
        # new_data[i, j] = 1 表示节点 i 有标签 j。
        for i,x in enumerate(data):
            if x is not None:
                for j in x:
                    new_data[i, j] = 1
        # 将 new_data 赋值给 labels['data']。
        labels['data'] = new_data
        # 将 mask 赋值给 labels['mask']。
        labels['mask'] = mask
        # 返回 labels 字典。
        return labels

    # 获取节点类型
    def get_node_type(self, node_id):
        for i in range(len(self.nodes['shift'])):
            if node_id < self.nodes['shift'][i]+self.nodes['count'][i]:
                return i

    # 获取边类型
    def get_edge_type(self, info):
        if type(info) is int or len(info) == 1:
            return info
        # 正向边
        for i in range(len(self.links['meta'])):
            if self.links['meta'][i] == info:
                return i
        # 如果不是正向边，则是反向边，则边类型<0
        info = (info[1], info[0])
        for i in range(len(self.links['meta'])):
            if self.links['meta'][i] == info:
                return -i - 1
        raise Exception('No available edge type')

    # 获取边的信息：由哪两种节点类型组成
    def get_edge_info(self, edge_id):
        return self.links['meta'][edge_id]

    # 将对应的链接数据列表转换为稀疏矩阵
    def list_to_sp_mat(self, li):
        data = [x[2] for x in li]
        i = [x[0] for x in li]
        j = [x[1] for x in li]
        return sp.coo_matrix((data, (i,j)), shape=(self.nodes['total'], self.nodes['total'])).tocsr()

    # 定义函数 load_links。
    # 函数的文档字符串描述了返回的 links 字典的结构，包括 total、count、meta 和 data 四个部分。
    def load_links(self):
        """
        return links dict
            total: total number of links
            count: a dict of int, number of links for each type
            meta: a dict of tuple, explaining the link type is from what type of node to what type of node
            data: a dict of sparse matrices, each link type with one matrix. Shapes are all (nodes['total'], nodes['total'])
        """
        # 初始化一个空的 links 字典，用于存储所有的链接信息：
        # total：链接总数，初始值为 0。
        # count：每种链接类型的数量，使用 Counter 进行统计。
        # meta：每种链接类型的元数据，存储链接类型对应的节点类型对。
        # data：每种链接类型的具体链接数据，使用 defaultdict(list) 存储。
        # todo data不知道是什么类型的数据，待检查
        links = {'total':0, 'count':Counter(), 'meta':{}, 'data':defaultdict(list)}
        # 打开链接数据文件 link.dat，文件路径是 self.path 和 link.dat 的拼接，读取模式为只读（'r'），编码为 'utf-8'。
        with open(os.path.join(self.path, 'link.dat'), 'r', encoding='utf-8') as f:
            for line in f:
                # 逐行读取文件内容，将每一行按照制表符 \t 进行分割。
                # 将每一行的字段转换为适当的数据类型：源节点 ID (h_id)、目标节点 ID (t_id)、关系类型 ID (r_id)、链接权重 (link_weight)，分别转换为整数和浮点数。
                th = line.split('\t')
                h_id, t_id, r_id, link_weight = int(th[0]), int(th[1]), int(th[2]), float(th[3])
                # 检查当前链接类型 r_id 是否在 links['meta'] 中：
                # 如果不在，则调用 get_node_type 函数获取源节点和目标节点的类型，存储在 links['meta'][r_id] 中。
                if r_id not in links['meta']:
                    h_type = self.get_node_type(h_id)
                    t_type = self.get_node_type(t_id)
                    links['meta'][r_id] = (h_type, t_type)
                # 将链接数据 (h_id, t_id, link_weight) 添加到 links['data'][r_id] 中。
                # 更新链接类型 r_id 的计数 links['count'][r_id]，加 1。
                # 更新链接总数 links['total']，加 1。
                links['data'][r_id].append((h_id, t_id, link_weight))
                links['count'][r_id] += 1
                links['total'] += 1
        # 初始化 new_data 空字典。
        # 遍历 links['data'] 中的每种链接类型 r_id，将其对应的链接数据列表转换为稀疏矩阵，存储在 new_data[r_id] 中。
        # 最后，将 new_data 赋值给 links['data']。
        new_data = {}
        for r_id in links['data']:
            new_data[r_id] = self.list_to_sp_mat(links['data'][r_id])
        links['data'] = new_data
        # 返回包含所有链接信息的 links 字典。
        return links

    # load_nodes 函数加载节点数据，返回包含节点信息的字典。
    # 这个方法主要用于从文件中读取节点数据，并构建一个包含节点总数、节点类型计数、节点属性和节点起始索引的字典。
    def load_nodes(self):
        """
        return nodes dict
            total: total number of nodes
            count: a dict of int, number of nodes for each type
            attr: a dict of np.array (or None), attribute matrices for each type of nodes
            shift: node_id shift for each type. You can get the id range of a type by
                        [ shift[node_type], shift[node_type]+count[node_type] )
        """
        # 初始化 nodes 字典，其中包含：
        # total: 节点总数，初始值为 0。
        # count: 记录每种类型节点的计数，使用 Counter 来统计。
        # attr: 存储节点属性的字典。
        # shift: 记录每种类型节点的起始索引。
        # 数据类型可能是下面这样（每个节点类型的节点扎堆）：
        # 节点ID    节点类型    属性
        # 0         0          [1.0, 2.0]
        # 1         0          [2.0, 3.0]
        # 2         0          [3.0, 4.0]
        # 3         0          [4.0, 5.0]
        # 4         1          无属性
        # 5         1          无属性
        # 6         1          无属性
        # 7         1          无属性
        # 8         1          无属性
        # 9         1          无属性
        nodes = {'total':0, 'count':Counter(), 'attr':{}, 'shift':{}}
        # 打开名为 'node.dat' 的文件，路径由 self.path 指定。以只读模式 ('r') 打开，并使用 UTF-8 编码。
        with open(os.path.join(self.path, 'node.dat'), 'r', encoding='utf-8') as f:
            # 遍历文件中的每一行，使用制表符 ('\t') 分割行内容，将其存储在列表 th 中。
            for line in f:
                th = line.split('\t')
                # 如果 th 列表的长度为 4，说明这一行包含完整的节点信息。
                # node_id: 节点 ID，将其转换为整数。
                # node_name: 节点名称（未使用）。
                # node_type: 节点类型，将其转换为整数。
                # node_attr: 节点属性，将属性字符串按逗号分割并转换为浮点数列表。
                if len(th) == 4:
                    # Then this line of node has attribute
                    node_id, node_name, node_type, node_attr = th
                    node_id = int(node_id)
                    node_type = int(node_type)
                    node_attr = list(map(float, node_attr.split(',')))
                    # 更新 nodes 字典：
                    # nodes['count'][node_type]: 增加节点类型的计数。
                    # nodes['attr'][node_id]: 存储节点属性。
                    # nodes['total']: 增加节点总数。
                    nodes['count'][node_type] += 1
                    nodes['attr'][node_id] = node_attr
                    nodes['total'] += 1
                # 如果 th 列表的长度为 3，说明这一行不包含属性信息。
                # node_id: 节点 ID，将其转换为整数。
                # node_name: 节点名称（未使用）。
                # node_type: 节点类型，将其转换为整数。
                elif len(th) == 3:
                    # Then this line of node doesn't have attribute
                    node_id, node_name, node_type = th
                    node_id = int(node_id)
                    node_type = int(node_type)
                    # 更新 nodes 字典：
                    # nodes['count'][node_type]: 增加节点类型的计数。
                    # nodes['total']: 增加节点总数。
                    nodes['count'][node_type] += 1
                    nodes['total'] += 1
                # 如果 th 列表的长度既不是 3 也不是 4，抛出异常，表示该行数据格式不正确。
                else:
                    raise Exception("Too few information to parse!")
        # 初始化变量：
        # shift: 用于记录每种类型节点的起始索引，初始值为 0。
        # attr: 用于存储每种类型节点的属性矩阵。
        shift = 0
        attr = {}
        # 遍历所有节点类型，为每种类型记录起始索引：
        # nodes['shift'][i]: 记录当前节点类型的起始索引。
        for i in range(len(nodes['count'])):
            nodes['shift'][i] = shift
            # 如果 shift 位置在 nodes['attr'] 中存在属性：
            # 遍历从 shift 到 shift + nodes['count'][i] 范围内的节点，收集属性到 mat 列表中。
            # 将 mat 转换为 NumPy 数组，并存储到 attr 字典中。
            if shift in nodes['attr']:
                mat = []
                for j in range(shift, shift+nodes['count'][i]):
                    mat.append(nodes['attr'][j])
                attr[i] = np.array(mat)
            # 如果 shift 位置在 nodes['attr'] 中不存在属性，将该节点类型的属性设置为 None。
            else:
                attr[i] = None
            # 更新 shift 变量，将其增加当前类型节点的数量，以便下一个类型的节点使用新的起始索引。
            shift += nodes['count'][i]
        # 将最终的属性字典 attr 赋值给 nodes['attr']，并返回 nodes 字典。
        nodes['attr'] = attr
        return nodes